use std::collections::HashMap;
use std::io::{BufReader, BufWriter, Write as _};
use std::net::{TcpListener, TcpStream, ToSocketAddrs};

use crossbeam::channel::{Receiver, Sender};
use log::{debug, error, info};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::encoding::{self, Value as _};
use crate::error::Result;
use crate::raft;
use crate::sql;
use crate::sql::engine::{Catalog as _, Engine as _};
use crate::sql::execution::StatementResult;
use crate::sql::types::{Row, Table};
use crate::storage;

/// A toyDB server. Routes messages to/from an inner Raft node.
///
/// * Listens for inbound SQL connections from clients via TCP and passes
///   requests to the local Raft node.
///
/// * Listens for inbound Raft connections from other toyDB nodes via TCP and
///   passes messages to the local Raft node.
///
/// * Connects to other toyDB nodes via TCP and sends outbound Raft messages
///   from the local Raft node.
pub struct Server {
    /// The inner Raft node.
    node: raft::Node,
    /// Outbound messages from the Raft node.
    node_rx: Receiver<raft::Envelope>,
}

impl Server {
    /// Creates a new toyDB server.
    pub fn new(
        id: raft::NodeID,
        raft_log: raft::Log,
        raft_state: Box<dyn raft::State>,
    ) -> Result<Self> {
        let (node_tx, node_rx) = crossbeam::channel::unbounded();
        let node = raft::Node::new(
            id,
            raft_log,
            raft_state,
            node_tx,
        )?;
        Ok(Self { node, node_rx })
    }

    /// Serves Raft and SQL requests indefinitely. Consumes the server.
    pub fn serve(self, sql_addr: impl ToSocketAddrs) -> Result<()> {
        let sql_listener = TcpListener::bind(sql_addr)?;
        info!(
            "Listening on {} (SQL) ",
            sql_listener.local_addr()?,
        );

        std::thread::scope(move |s| {
            let id = self.node.id();
            let (raft_request_tx, raft_request_rx) = crossbeam::channel::unbounded();

            // Route Raft messages between the local node, peers, and clients.
            s.spawn(move || {
                Self::raft_route(
                    self.node,
                    self.node_rx,
                    raft_request_rx,
                )
            });

            // Serve inbound SQL connections.
            let sql_engine = sql::engine::Raft::new(raft_request_tx);
            s.spawn(move || Self::sql_accept(id, sql_listener, sql_engine));
        });

        Ok(())
    }


    /// Routes Raft messages:
    ///
    /// * node_rx: outbound messages from the local Raft node. Routed to peers
    ///   via TCP, or to local clients via a response channel.
    ///
    /// * request_rx: inbound requests from local SQL clients. Stepped into
    ///   the local Raft node as ClientRequest messages. Responses are returned
    ///   via the provided response channel.
    ///
    /// * peers_rx: inbound messages from remote Raft peers. Stepped into the
    ///   local Raft node.
    ///
    /// * peers_tx: outbound per-peer channels sent via TCP connections.
    ///   Messages from the local node's node_rx are sent here.
    ///
    /// Panics on any errors, since the Raft node can't recover from failed
    /// state transitions.
    fn raft_route(
        mut node: raft::Node,
        node_rx: Receiver<raft::Envelope>,
        request_rx: Receiver<(raft::Request, Sender<Result<raft::Response>>)>,
    ) {
        // Track response channels by request ID. The Raft node will emit
        // ClientResponse messages that we forward to the response channel.
        let mut response_txs = HashMap::<raft::RequestID, Sender<Result<raft::Response>>>::new();

        loop {
            crossbeam::select! {


                // If we receive a client response addressed to the local node,
                // forward it to the waiting client via the response channel.
                recv(node_rx) -> result => {
                    let msg = result.expect("node_rx disconnected");
                    if msg.to == node.id()
                        && let raft::Message::ClientResponse{ id, response } = msg.message
                    {
                        if let Some(response_tx) = response_txs.remove(&id) {
                            response_tx.send(response).expect("response_tx disconnected");
                        }
                        continue
                    }
                }

                // Track inbound client requests and step them into the node.
                recv(request_rx) -> result => {
                    let (request, response_tx) = result.expect("request_rx disconnected");
                    let id = Uuid::new_v4();
                    let msg = raft::Envelope{
                        from: node.id(),
                        to: node.id(),
                        term: node.term(),
                        message: raft::Message::ClientRequest{id, request},
                    };
                    node = node.step(msg).expect("step failed");
                    response_txs.insert(id, response_tx);
                }
            }
        }
    }

    /// Accepts new SQL client connections and spawns session threads for them.
    fn sql_accept(id: raft::NodeID, listener: TcpListener, sql_engine: sql::engine::Raft) {
        std::thread::scope(|s| {
            loop {
                let (socket, peer) = match listener.accept() {
                    Ok((socket, peer)) => (socket, peer),
                    Err(err) => {
                        error!("Client accept failed: {err}");
                        continue;
                    }
                };
                let session = sql_engine.session();
                s.spawn(move || {
                    debug!("Client {peer} connected");
                    match Self::sql_session(id, socket, session) {
                        Ok(()) => debug!("Client {peer} disconnected"),
                        Err(err) => error!("Client {peer} error: {err}"),
                    }
                });
            }
        })
    }

    /// Processes a client SQL session, executing SQL statements against the
    /// Raft node.
    fn sql_session(
        id: raft::NodeID,
        socket: TcpStream,
        mut session: sql::execution::Session<sql::engine::Raft>,
    ) -> Result<()> {
        let mut reader = BufReader::new(socket.try_clone()?);
        let mut writer = BufWriter::new(socket);

        while let Some(request) = Request::maybe_decode_from(&mut reader)? {
            // Execute request.
            debug!("Received request {request:?}");
            let response = match request {
                Request::Execute(query) => session.execute(&query).map(Response::Execute),
                Request::GetTable(table) => {
                    session.with_txn(true, |txn| txn.must_get_table(&table)).map(Response::GetTable)
                }
                Request::ListTables => session
                    .with_txn(true, |txn| {
                        Ok(txn.list_tables()?.into_iter().map(|t| t.name).collect())
                    })
                    .map(Response::ListTables),
                Request::Status => session
                    .status()
                    .map(|s| Status { server: id, raft: s.raft, mvcc: s.mvcc })
                    .map(Response::Status),
            };

            // Process response.
            debug!("Returning response {response:?}");
            response.encode_into(&mut writer)?;
            writer.flush()?;
        }
        Ok(())
    }
}

/// A SQL client request.
#[derive(Debug, Serialize, Deserialize)]
pub enum Request {
    /// Executes a SQL statement.
    Execute(String),
    /// Fetches the given table schema.
    GetTable(String),
    /// Lists all tables.
    ListTables,
    /// Returns server status.
    Status,
}

impl encoding::Value for Request {}

/// A SQL server response.
#[derive(Debug, Serialize, Deserialize)]
pub enum Response {
    Execute(StatementResult),
    Row(Option<Row>),
    GetTable(Table),
    ListTables(Vec<String>),
    Status(Status),
}

impl encoding::Value for Response {}

/// SQL server status.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Status {
    pub server: raft::NodeID,
    pub raft: raft::Status,
    pub mvcc: storage::mvcc::Status,
}
